{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "KUu4vOt5zI9d"
   },
   "outputs": [],
   "source": [
    "# Copyright 2018 The TensorFlow Hub Authors. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "edfbxDDh2AEs"
   },
   "source": [
    "## Predict Shakespeare with Cloud TPUs and Keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RNo1Vfghpa8j"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This example uses [tf.keras](https://www.tensorflow.org/guide/keras) to build a *language model* and train it on a Cloud TPU. This language model predicts the next character of text given the text so far. The trained model can generate new snippets of text that read in a similar style to the text training data.\n",
    "\n",
    "The model trains for 10 epochs and completes in approximately 5 minutes.\n",
    "\n",
    "This notebook is hosted on GitHub. To view it in its original repository, after opening the notebook in Colab, select **File > View on GitHub**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dgAHfQtuhddd"
   },
   "source": [
    "## Learning objectives\n",
    "\n",
    "In this notebook, you will learn how to:\n",
    "*   Build a two-layer, forward-LSTM model.\n",
    "*   Convert a `tf.keras` model to an equivalent TPU version and then use the standard Keras methods to train: `fit`, `predict`, and `evaluate`.\n",
    "*   Use the trained model to make predictions and generate your own Shakespeare-esque play.\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instruction to Train on Deep Learning VM\n",
    "\n",
    "* Start a VM attached to a TPU:\n",
    "```\n",
    "    INSTANCE_NAME=mytpu   # CHANGE THIS\n",
    "    GCP_LOGIN_NAME=google-cloud-customer@gmail.com  # CHANGE THIS\n",
    "    TPU_NAME=$INSTANCE_NAME\n",
    "    \n",
    "    gcloud compute instances create $INSTANCE_NAME \\\n",
    "    --machine-type n1-standard-8 \\\n",
    "    --image-project deeplearning-platform-release \\\n",
    "    --image-family tf-1-12-cpu \\\n",
    "    --scopes cloud-platform \\\n",
    "    --metadata proxy-user-mail=\"${GCP_LOGIN_NAME}\",\\\n",
    "    startup-script=\"echo export TPU_NAME=$TPU_NAME > /etc/profile.d/tpu-env.sh\"\n",
    "    \n",
    "    gcloud compute tpus create $TPU_NAME \\\n",
    "     --network default \\\n",
    "     --range 10.240.1.0 \\\n",
    "     --version 1.12\n",
    "```\n",
    "* Open JupyterLab from the GCP web console and clone the repo that this notebook is hosted in:\n",
    "   ```\n",
    "       git clone https://www.github.com/GoogleCloudPlatform/training-data-analyst\n",
    "   ```\n",
    "   and navigate to courses/fast-and-lean-data-science/09_lstm_keras_tpu.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_I0RdnOSkNmi"
   },
   "source": [
    "#### Alternate method: Train on Colab\n",
    "\n",
    "   1. On the main menu, click Runtime and select **Change runtime type**. Set \"TPU\" as the hardware accelerator.\n",
    "   1. Click Runtime again and select **Runtime > Run All**. You can also run the cells manually with Shift-ENTER. \n",
    "\n",
    "Note that the Colab method is suitable for educational use, but you will not be able to use it for long-term work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "SEQ_LEN = 100\n",
    "BATCH_SIZE = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Lvo0t7XVIkWZ"
   },
   "source": [
    "## Data, model, and training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xzpUtDMqmA-x"
   },
   "source": [
    "In this example, you train the model on the combined works of William Shakespeare, then use the model to compose a play in the style of *The Great Bard*:\n",
    "\n",
    "<blockquote>\n",
    "Loves that led me no dumbs lack her Berjoy's face with her to-day.  \n",
    "The spirits roar'd; which shames which within his powers  \n",
    "\tWhich tied up remedies lending with occasion,  \n",
    "A loud and Lancaster, stabb'd in me  \n",
    "\tUpon my sword for ever: 'Agripo'er, his days let me free.  \n",
    "\tStop it of that word, be so: at Lear,  \n",
    "\tWhen I did profess the hour-stranger for my life,  \n",
    "\tWhen I did sink to be cried how for aught;  \n",
    "\tSome beds which seeks chaste senses prove burning;  \n",
    "But he perforces seen in her eyes so fast;  \n",
    "And _  \n",
    "</blockquote>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KRQ6Fjra3Ruq"
   },
   "source": [
    "### Data\n",
    "\n",
    "We have downloaded *The Complete Works of William Shakespeare* as a single text file from [Project Gutenberg](https://www.gutenberg.org/) and stored in a GCS bucket. Because TPUs are located in Google Cloud, for optimal performance, they read data directly from Google Cloud Storage (GCS).\n",
    "\n",
    "We will use snippets from this file as the *training data* for the model. The *target* snippet is offset by one character."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "From hands of falsehood, in sure wards of trust!\n",
      "But thou, to whom my jewels trifles are,\n",
      "Most worthy comfort, now my greatest grief,\n",
      "Thou best of dearest, and mine only care,\n",
      "Art left the prey of every vulgar thief.\n",
      "Thee have I not locked up in any chest,\n",
      "Save where thou art not, though I feel thou art,\n",
      "Within the gentle closure of my breast,\n",
      "From whence at pleasure thou mayst come and part,\n",
      "  And even thence thou wilt be stol’n I fear,\n"
     ]
    }
   ],
   "source": [
"!gcloud storage cat gs://cloud-training-demos/tpudemo/shakespeare.txt | head -1000 | tail -10"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AbL6cqCl7hnt"
   },
   "source": [
    "### Build the tf.data.Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "E3V4V-Jxmuv3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Input text [5796379] ﻿\n",
      "Project Gutenberg’s The Complete Works of Willi\n",
      "(array([[115,  32, 116, 111,  32,  97,  32,  98,  97, 119]], dtype=int32), array([[[ 32],\n",
      "        [116],\n",
      "        [111],\n",
      "        [ 32],\n",
      "        [ 97],\n",
      "        [ 32],\n",
      "        [ 98],\n",
      "        [ 97],\n",
      "        [119],\n",
      "        [100]]], dtype=int32))\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import six\n",
    "import tensorflow as tf\n",
    "import time\n",
    "import os\n",
    "\n",
    "SHAKESPEARE_TXT = 'gs://cloud-training-demos/tpudemo/shakespeare.txt'\n",
    "\n",
    "tf.logging.set_verbosity(tf.logging.INFO)\n",
    "\n",
    "def transform(txt, pad_to=None):\n",
    "  # drop any non-ascii characters\n",
    "  output = np.asarray([ord(c) for c in txt if ord(c) < 255], dtype=np.int32)\n",
    "  if pad_to is not None:\n",
    "    output = output[:pad_to]\n",
    "    output = np.concatenate([\n",
    "        np.zeros([pad_to - len(txt)], dtype=np.int32),\n",
    "        output,\n",
    "    ])\n",
    "  return output\n",
    "\n",
    "def training_generator(seq_len=SEQ_LEN, batch_size=BATCH_SIZE):\n",
    "  \"\"\"A generator yields (source, target) arrays for training.\"\"\"\n",
    "  with tf.gfile.GFile(SHAKESPEARE_TXT, 'r') as f:\n",
    "    txt = f.read()\n",
    "\n",
    "  tf.logging.info('Input text [%d] %s', len(txt), txt[:50])\n",
    "  source = transform(txt)\n",
    "  while True:\n",
    "    offsets = np.random.randint(0, len(source) - seq_len, batch_size)\n",
    "\n",
    "    # Our model uses sparse crossentropy loss, but Keras requires labels\n",
    "    # to have the same rank as the input logits.  We add an empty final\n",
    "    # dimension to account for this.\n",
    "    yield (\n",
    "        np.stack([source[idx:idx + seq_len] for idx in offsets]),\n",
    "        np.expand_dims(\n",
    "            np.stack([source[idx + 1:idx + seq_len + 1] for idx in offsets]),\n",
    "            -1),\n",
    "    )\n",
    "\n",
    "a = six.next(training_generator(seq_len=10, batch_size=1))\n",
    "print(a)\n",
    "#print(tf.convert_to_tensor(a[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset():\n",
    "    return tf.data.Dataset.from_generator(training_generator, \n",
    "                                          (tf.int32, tf.int32),\n",
    "                                          (tf.TensorShape([BATCH_SIZE, SEQ_LEN]), \n",
    "                                           tf.TensorShape([BATCH_SIZE, SEQ_LEN, 1]))\n",
    "                                         )\n",
    "                                          "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Bbb05dNynDrQ"
   },
   "source": [
    "### Build the model\n",
    "\n",
    "The model is defined as a two-layer, forward-LSTM—with two changes from the `tf.keras` standard LSTM definition:\n",
    "\n",
    "1. Define the input `shape` of the model to comply with the [XLA compiler](https://www.tensorflow.org/performance/xla/)'s static shape requirement.\n",
    "2. Use `tf.train.Optimizer` instead of a standard Keras optimizer (Keras optimizer support is still experimental)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yLEM-fLJlEEt"
   },
   "outputs": [],
   "source": [
    "EMBEDDING_DIM = 512\n",
    "\n",
    "def lstm_model(seq_len, batch_size, stateful):\n",
    "  \"\"\"Language model: predict the next word given the current word.\"\"\"\n",
    "  source = tf.keras.Input(\n",
    "      name='seed', shape=(seq_len,), batch_size=batch_size, dtype=tf.int32)\n",
    "\n",
    "  embedding = tf.keras.layers.Embedding(input_dim=256, output_dim=EMBEDDING_DIM)(source)\n",
    "  lstm_1 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(embedding)\n",
    "  lstm_2 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(lstm_1)\n",
    "  predicted_char = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(256, activation='softmax'))(lstm_2)\n",
    "  model = tf.keras.Model(inputs=[source], outputs=[predicted_char])\n",
    "  model.compile(\n",
    "      optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01),\n",
    "      loss='sparse_categorical_crossentropy',\n",
    "      metrics=['sparse_categorical_accuracy'])\n",
    "  return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VzBYDJI0_Tfm"
   },
   "source": [
    "### Train the model\n",
    "\n",
    "The `tf.contrib.tpu.keras_to_tpu_model` function converts a `tf.keras` model to an equivalent TPU version. You then use the standard Keras methods to train: `fit`, `predict`, and `evaluate`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ExQ922tfzSGA"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Querying Tensorflow master (grpc://10.240.1.2:8470) for TPU system metadata.\n",
      "INFO:tensorflow:Found TPU system:\n",
      "INFO:tensorflow:*** Num TPU Cores: 8\n",
      "INFO:tensorflow:*** Num TPU Workers: 1\n",
      "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, 3577263294212491411)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 9025790428816496025)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_GPU:0, XLA_GPU, 17179869184, 12633420940520028490)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 4889055515393833545)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 9632346120050048346)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 17116235810456753864)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 10387310604824546591)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 1907352312443083400)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 12795733773701058027)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 3885015511603421288)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 5505781315448043335)\n",
      "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 4880015139463075197)\n",
      "WARNING:tensorflow:tpu_model (from tensorflow.contrib.tpu.python.tpu.keras_support) is experimental and may change or be removed at any time, and without warning.\n",
      "Running on TPU  ['10.240.1.2:8470']\n",
      "Epoch 1/10\n",
      "INFO:tensorflow:New input shapes; (re-)compiling: mode=train (# of cores 8), [TensorSpec(shape=(128,), dtype=tf.int32, name=None), TensorSpec(shape=(128, 100), dtype=tf.int32, name=None), TensorSpec(shape=(128, 100, 1), dtype=tf.int32, name=None)]\n",
      "INFO:tensorflow:Overriding default placeholder.\n",
      "INFO:tensorflow:Remapping placeholder for seed\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "Input 'y' of 'Equal' Op has type int64 that does not match type int32 of argument 'x'.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/op_def_library.py\u001b[0m in \u001b[0;36m_apply_op_helper\u001b[0;34m(self, op_type_name, name, **keywords)\u001b[0m\n\u001b[1;32m    509\u001b[0m                 \u001b[0mas_ref\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minput_arg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_ref\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 510\u001b[0;31m                 preferred_dtype=default_dtype)\n\u001b[0m\u001b[1;32m    511\u001b[0m           \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0merr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/ops.py\u001b[0m in \u001b[0;36minternal_convert_to_tensor\u001b[0;34m(value, dtype, name, as_ref, preferred_dtype, ctx)\u001b[0m\n\u001b[1;32m   1145\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mret\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1146\u001b[0;31m       \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconversion_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mas_ref\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mas_ref\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1147\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/ops.py\u001b[0m in \u001b[0;36m_TensorTensorConversionFunction\u001b[0;34m(t, dtype, name, as_ref)\u001b[0m\n\u001b[1;32m    982\u001b[0m         \u001b[0;34m\"Tensor conversion requested dtype %s for Tensor with dtype %s: %r\"\u001b[0m \u001b[0;34m%\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 983\u001b[0;31m         (dtype.name, t.dtype.name, str(t)))\n\u001b[0m\u001b[1;32m    984\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Tensor conversion requested dtype int32 for Tensor with dtype int64: 'Tensor(\"tpu_139640912195936/metrics/sparse_categorical_accuracy/ArgMax:0\", shape=(1024, 100), dtype=int64)'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-fbd398209b08>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0mtraining_input\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m     \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m10\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m )\n\u001b[1;32m     24\u001b[0m \u001b[0mtpu_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave_weights\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'/tmp/bard.h5'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moverwrite\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, **kwargs)\u001b[0m\n\u001b[1;32m   1505\u001b[0m                                   \u001b[0mvalidation_split\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalidation_data\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshuffle\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1506\u001b[0m                                   \u001b[0mclass_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minitial_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1507\u001b[0;31m                                   steps_per_epoch, validation_steps, **kwargs)\n\u001b[0m\u001b[1;32m   1508\u001b[0m       \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1509\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_numpy_to_infeed_manager_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36m_pipeline_fit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, **kwargs)\u001b[0m\n\u001b[1;32m   1598\u001b[0m         \u001b[0minitial_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minitial_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1599\u001b[0m         \u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1600\u001b[0;31m         validation_steps=validation_steps)\n\u001b[0m\u001b[1;32m   1601\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1602\u001b[0m   def _pipeline_fit_loop(self,\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36m_pipeline_fit_loop\u001b[0;34m(self, inputs, targets, sample_weights, batch_size, epochs, verbose, callbacks, val_inputs, val_targets, val_sample_weights, shuffle, initial_epoch, steps_per_epoch, validation_steps)\u001b[0m\n\u001b[1;32m   1685\u001b[0m             \u001b[0mval_sample_weights\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mval_sample_weights\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1686\u001b[0m             \u001b[0mvalidation_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mvalidation_steps\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1687\u001b[0;31m             epoch_logs=epoch_logs)\n\u001b[0m\u001b[1;32m   1688\u001b[0m       \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1689\u001b[0m         \u001b[0;31m# Sample-wise fit loop.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36m_pipeline_fit_loop_step_wise\u001b[0;34m(self, ins, callbacks, steps_per_epoch, epochs, do_validation, val_inputs, val_targets, val_sample_weights, validation_steps, epoch_logs)\u001b[0m\n\u001b[1;32m   1813\u001b[0m     \u001b[0;31m# Loop prologue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1814\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1815\u001b[0;31m       \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpipeline_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcur_step_inputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnext_step_inputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mins\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1816\u001b[0m       \u001b[0;32massert\u001b[0m \u001b[0mouts\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m  \u001b[0;31m# Function shouldn't return anything!\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1817\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0merrors\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mOutOfRangeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36mpipeline_run\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m   1323\u001b[0m           next_input_tensors)\n\u001b[1;32m   1324\u001b[0m       next_tpu_model_ops = self._tpu_model_ops_for_input_specs(\n\u001b[0;32m-> 1325\u001b[0;31m           next_input_specs, next_step_infeed_manager)\n\u001b[0m\u001b[1;32m   1326\u001b[0m       \u001b[0minfeed_dict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnext_infeed_instance\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmake_feed_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext_tpu_model_ops\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1327\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36m_tpu_model_ops_for_input_specs\u001b[0;34m(self, input_specs, infeed_manager)\u001b[0m\n\u001b[1;32m   1152\u001b[0m           self._tpu_assignment.num_towers, input_specs)\n\u001b[1;32m   1153\u001b[0m       new_tpu_model_ops = self._specialize_model(input_specs,\n\u001b[0;32m-> 1154\u001b[0;31m                                                  infeed_manager)\n\u001b[0m\u001b[1;32m   1155\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compilation_cache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mshape_key\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnew_tpu_model_ops\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1156\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_test_model_compiles\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnew_tpu_model_ops\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36m_specialize_model\u001b[0;34m(self, input_specs, infeed_manager)\u001b[0m\n\u001b[1;32m   1062\u001b[0m     \u001b[0;31m# running on a different logical core.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1063\u001b[0m     compile_op, execute_op = tpu.split_compile_and_replicate(\n\u001b[0;32m-> 1064\u001b[0;31m         _model_fn, inputs=[[]] * self._tpu_assignment.num_towers)\n\u001b[0m\u001b[1;32m   1065\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1066\u001b[0m     \u001b[0;31m# Generate CPU side operations to enqueue features/labels and dequeue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/tpu.py\u001b[0m in \u001b[0;36msplit_compile_and_replicate\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    682\u001b[0m       \u001b[0mvscope\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_custom_getter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcustom_getter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    683\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 684\u001b[0;31m       \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcomputation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mcomputation_inputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    685\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    686\u001b[0m       \u001b[0mvscope\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_use_resource\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msaved_use_resource\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/contrib/tpu/python/tpu/keras_support.py\u001b[0m in \u001b[0;36m_model_fn\u001b[0;34m()\u001b[0m\n\u001b[1;32m   1003\u001b[0m                   weighted_metrics=metrics_module.clone_metrics(\n\u001b[1;32m   1004\u001b[0m                       self.model.weighted_metrics),\n\u001b[0;32m-> 1005\u001b[0;31m                   \u001b[0mtarget_tensors\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtpu_targets\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1006\u001b[0m               )\n\u001b[1;32m   1007\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/training/checkpointable/base.py\u001b[0m in \u001b[0;36m_method_wrapper\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    472\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_setattr_tracking\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    473\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 474\u001b[0;31m       \u001b[0mmethod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    475\u001b[0m     \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    476\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_setattr_tracking\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprevious_value\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mcompile\u001b[0;34m(self, optimizer, loss, metrics, loss_weights, sample_weight_mode, weighted_metrics, target_tensors, distribute, **kwargs)\u001b[0m\n\u001b[1;32m    646\u001b[0m         \u001b[0mtargets\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtargets\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    647\u001b[0m         \u001b[0mskip_target_indices\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mskip_target_indices\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 648\u001b[0;31m         sample_weights=self.sample_weights)\n\u001b[0m\u001b[1;32m    649\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    650\u001b[0m     \u001b[0;31m# Prepare gradient updates and state updates.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36m_handle_metrics\u001b[0;34m(self, outputs, skip_target_indices, targets, sample_weights, masks)\u001b[0m\n\u001b[1;32m    311\u001b[0m         metric_results.extend(\n\u001b[1;32m    312\u001b[0m             self._handle_per_output_metrics(self._per_output_metrics[i], target,\n\u001b[0;32m--> 313\u001b[0;31m                                             output, output_mask))\n\u001b[0m\u001b[1;32m    314\u001b[0m         metric_results.extend(\n\u001b[1;32m    315\u001b[0m             self._handle_per_output_metrics(\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36m_handle_per_output_metrics\u001b[0;34m(self, metrics_dict, y_true, y_pred, mask, weights)\u001b[0m\n\u001b[1;32m    268\u001b[0m               metric_fn)\n\u001b[1;32m    269\u001b[0m           metric_result = weighted_metric_fn(\n\u001b[0;32m--> 270\u001b[0;31m               y_true, y_pred, weights=weights, mask=mask)\n\u001b[0m\u001b[1;32m    271\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    272\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecuting_eagerly\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/engine/training_utils.py\u001b[0m in \u001b[0;36mweighted\u001b[0;34m(y_true, y_pred, weights, mask)\u001b[0m\n\u001b[1;32m    596\u001b[0m     \"\"\"\n\u001b[1;32m    597\u001b[0m     \u001b[0;31m# score_array has ndim >= 2\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 598\u001b[0;31m     \u001b[0mscore_array\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_true\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_pred\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    599\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mmask\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    600\u001b[0m       \u001b[0mmask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmath_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_pred\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/keras/metrics.py\u001b[0m in \u001b[0;36msparse_categorical_accuracy\u001b[0;34m(y_true, y_pred)\u001b[0m\n\u001b[1;32m    660\u001b[0m     \u001b[0my_pred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmath_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_pred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mK\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloatx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    661\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 662\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mmath_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmath_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mequal\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_true\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_pred\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mK\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloatx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    663\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    664\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/gen_math_ops.py\u001b[0m in \u001b[0;36mequal\u001b[0;34m(x, y, name)\u001b[0m\n\u001b[1;32m   2732\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0m_ctx\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0m_ctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_eager_context\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_eager\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2733\u001b[0m     _, _, _op = _op_def_lib._apply_op_helper(\n\u001b[0;32m-> 2734\u001b[0;31m         \"Equal\", x=x, y=y, name=name)\n\u001b[0m\u001b[1;32m   2735\u001b[0m     \u001b[0m_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_op\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2736\u001b[0m     \u001b[0m_inputs_flat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_op\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/op_def_library.py\u001b[0m in \u001b[0;36m_apply_op_helper\u001b[0;34m(self, op_type_name, name, **keywords)\u001b[0m\n\u001b[1;32m    544\u001b[0m                   \u001b[0;34m\"%s type %s of argument '%s'.\"\u001b[0m \u001b[0;34m%\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    545\u001b[0m                   (prefix, dtypes.as_dtype(attrs[input_arg.type_attr]).name,\n\u001b[0;32m--> 546\u001b[0;31m                    inferred_from[input_arg.type_attr]))\n\u001b[0m\u001b[1;32m    547\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    548\u001b[0m           \u001b[0mtypes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: Input 'y' of 'Equal' Op has type int64 that does not match type int32 of argument 'x'."
     ]
    }
   ],
   "source": [
    "tf.keras.backend.clear_session()\n",
    "\n",
    "training_model = lstm_model(seq_len=SEQ_LEN, batch_size=BATCH_SIZE, stateful=False)\n",
    "\n",
    "# Use TPU if it exists, else fall back to GPU\n",
    "try: # TPU detection\n",
    "  tpu = tf.contrib.cluster_resolver.TPUClusterResolver()\n",
    "  training_model = tf.contrib.tpu.keras_to_tpu_model(\n",
    "        training_model,\n",
    "        strategy=tf.contrib.tpu.TPUDistributionStrategy(tpu))\n",
    "  training_input = create_dataset  # Function that returns a dataset\n",
    "  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n",
    "except ValueError:\n",
    "  tpu = None\n",
    "  training_input = create_dataset()  # The dataset itself\n",
    "  print(\"Running on GPU or CPU\")\n",
    "\n",
    "# Run fit()\n",
    "training_model.fit(\n",
    "    training_input,\n",
    "    steps_per_epoch=100,\n",
    "    epochs=10,\n",
    ")\n",
    "tpu_model.save_weights('/tmp/bard.h5', overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TCBtcpZkykSf"
   },
   "source": [
    "### Make predictions with the model\n",
    "\n",
    "Use the trained model to make predictions and generate your own Shakespeare-esque play.\n",
    "Start the model off with a *seed* sentence, then generate 250 characters from it. The model makes five predictions from the initial seed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tU7M-EGGxR3E"
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 5\n",
    "PREDICT_LEN = 250\n",
    "\n",
    "# Keras requires the batch size be specified ahead of time for stateful models.\n",
    "# We use a sequence length of 1, as we will be feeding in one character at a \n",
    "# time and predicting the next character.\n",
    "prediction_model = lstm_model(seq_len=1, batch_size=BATCH_SIZE, stateful=True)\n",
    "prediction_model.load_weights('/tmp/bard.h5')\n",
    "\n",
    "# We seed the model with our initial string, copied BATCH_SIZE times\n",
    "\n",
    "seed_txt = 'Looks it not like the king?  Verily, we must go! '\n",
    "seed = transform(seed_txt)\n",
    "seed = np.repeat(np.expand_dims(seed, 0), BATCH_SIZE, axis=0)\n",
    "\n",
    "# First, run the seed forward to prime the state of the model.\n",
    "prediction_model.reset_states()\n",
    "for i in range(len(seed_txt) - 1):\n",
    "  prediction_model.predict(seed[:, i:i + 1])\n",
    "\n",
    "# Now we can accumulate predictions!\n",
    "predictions = [seed[:, -1:]]\n",
    "for i in range(PREDICT_LEN):\n",
    "  last_word = predictions[-1]\n",
    "  next_probits = prediction_model.predict(last_word)[:, 0, :]\n",
    "  \n",
    "  # sample from our output distribution\n",
    "  next_idx = [\n",
    "      np.random.choice(256, p=next_probits[i])\n",
    "      for i in range(BATCH_SIZE)\n",
    "  ]\n",
    "  predictions.append(np.asarray(next_idx, dtype=np.int32))\n",
    "  \n",
    "\n",
    "for i in range(BATCH_SIZE):\n",
    "  print('PREDICTION %d\\n\\n' % i)\n",
    "  p = [predictions[j][i] for j in range(PREDICT_LEN)]\n",
    "  generated = ''.join([chr(c) for c in p])\n",
    "  print(generated)\n",
    "  print()\n",
    "  assert len(generated) == PREDICT_LEN, 'Generated text too short'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2a5cGsSTEBQD"
   },
   "source": [
    "## What's next\n",
    "\n",
    "* Learn about [Cloud TPUs](https://cloud.google.com/tpu/docs) that Google designed and optimized specifically to speed up and scale up ML workloads for training and inference and to enable ML engineers and researchers to iterate more quickly.\n",
    "* Explore the range of [Cloud TPU tutorials and Colabs](https://cloud.google.com/tpu/docs/tutorials) to find other examples that can be used when implementing your ML project.\n",
    "\n",
    "On Google Cloud Platform, in addition to GPUs and TPUs available on pre-configured [deep learning VMs](https://cloud.google.com/deep-learning-vm/),  you will find [AutoML](https://cloud.google.com/automl/)*(beta)* for training custom models without writing code and [Cloud ML Engine](https://cloud.google.com/ml-engine/docs/) which will allows you to run parallel trainings and hyperparameter tuning of your custom models on powerful distributed hardware.\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [
    "N6ZDpd9XzFeN"
   ],
   "name": "Predict Shakespeare with Cloud TPUs and Keras",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
